#include "DebugIntegrator.cuh"

using namespace OBR;

static __forceinline__ __device__ Spectrum customDebugShader(InteractionData &inter,
                                                             DebugIntegratorData &data)
{
    return data.default_val;
}

extern "C" __device__ Spectrum __continuation_callable__debug_integrator(SamplePRD *prd)
{
    const uint mat_base = prd->params->material_base_offset;
    const uint med_base = prd->params->medium_base_offset;
    const float eps = 1e-4;
    CallableData *sbt_record = reinterpret_cast<CallableData *>(optixGetSbtDataPointer());
    DebugIntegratorData *settings =
        reinterpret_cast<DebugIntegratorData *>(sbt_record->callable_payload);

    Spectrum debuginfo;
    uint u0, u1;
    packPointer(prd, u0, u1);

    // trace one ray to find 1st hit
    optixTrace(                        // trace new ray
        prd->params->handle,           // handel
        prd->inter.wi.o,               // ray origin
        prd->inter.wi.d,               // ray direction
        eps, 1e10f,                    // ray time
        prd->scene_time,               // scene time
        OptixVisibilityMask(1),        // vis mask
        OPTIX_RAY_FLAG_DISABLE_ANYHIT, // ray flags
        0,                             // SBT offset
        1,                             // SBT stride
        0,                             // miss sbt
        u0, u1);

    // handle hit case
    Spectrum tr = Spectrum(1.0f);
    if (prd->hit)
    {
        // handle medium interaction if wo is in medium
        if (prd->inter.wo.med_id != -1)
        {
            // if a medium interaction happens, prd->inter is updated, and has_mi is set to true
            tr = optixDirectCall<Spectrum, SamplePRD *, MediumCallMode>(
                prd->inter.wo.med_id + med_base, prd, MediumCallMode::TRACK_INTERACTION);
        }

        // short expression
        InteractionData &inter = prd->inter;

        switch (settings->mode)
        {
        case DebugMode::ALBEDO: {
            if (prd->inter.mat_id != -1)
            {
                float4 urand = curand_uniform4(prd->rng_long);
                float pdf;
                debuginfo = optixDirectCall<Spectrum, SamplePRD *, float *, float, float>(
                    prd->inter.mat_id + mat_base, prd, &pdf, urand.x, urand.y,
                    MaterialCallMode::SAMPLE_FR);
                debuginfo /= pdf;
            }
            else if (prd->inter.light_id != -1)
            {
                float pdf;
                debuginfo = optixDirectCall<Spectrum, SamplePRD *, float *, int, float, float,
                                            LightCallMode>(prd->inter.light_id + 1, prd, &pdf,
                                                           prd->inter.light_instance_id, 0, 0,
                                                           LightCallMode::CALC_LI_PDF);
                debuginfo /= pdf;
            }
            break;
        }
        case DebugMode::UV: {
            float debugcolor[3] = {inter.uv.x, inter.uv.y, 0.0f};
            debuginfo = debugcolor;
            break;
        }
        case DebugMode::NORMAL: {
            float3 n_ = inter.n * 0.5f + 0.5f;
            debuginfo = Spectrum((float *)&n_);
            break;
        }
        case DebugMode::DEPTH: {
            debuginfo = inter.wo.t * settings->scale;
            break;
        }
        case DebugMode::INTERACTION_TYPE: {
            bool is_si = (inter.type == InteractionType::SURFACE_INTERACTION);
            float debugcolor[3] = {0, 0, is_si ? 1 : 0};
            debuginfo = debugcolor;
            break;
        }
        case DebugMode::MAT_ID: {
            debuginfo = inter.mat_id * settings->scale;
            break;
        }
        case DebugMode::MAT_INSTANCE_ID: {
            debuginfo = inter.mat_id == -1 ? 0 : inter.mat_instance_id * settings->scale;
            break;
        }
        case DebugMode::LIGHT_ID: {
            debuginfo = inter.light_id * settings->scale;
            break;
        }
        case DebugMode::LIGHT_INSTANCE_ID: {
            debuginfo = inter.light_id == -1 ? 0 : inter.light_instance_id * settings->scale;
            break;
        }
        case DebugMode::MESH_ID: {
            debuginfo = inter.mesh_id * settings->scale;
            break;
        }
        case DebugMode::MESH_INSTANCE_ID: {
            debuginfo = inter.mesh_id == -1 ? 0 : inter.mesh_instance_id * settings->scale;
            break;
        }
        case DebugMode::TR: {
            debuginfo = tr;
            break;
        }
        case DebugMode::CUSTOM: {
            debuginfo = customDebugShader(inter, *settings);
            break;
        }
        default: {
            debuginfo = settings->default_val;
            break;
        }
        }
    }
    else
    {
        debuginfo = settings->default_val;
    }

    return debuginfo;
}